﻿using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.ComponentModel.DataAnnotations;
using System.Linq;
using System.Text;
using System.IO;
using System.Diagnostics;
using System.IO.MemoryMappedFiles;
using System.Text.RegularExpressions;
using System.Windows.Markup;
using System.Threading;

using System.Reactive.Linq;
using System.Reactive.Concurrency;
using System.Reactive.Threading.Tasks;
using System.Reactive.Linq.Observαble;
using System.Threading.Tasks;

using alib.Enumerable;
using alib.Character;
using alib.Dictionary;
using alib.Observable;

using agree.schema;

namespace agree
{
	public sealed class MaxEntModelLoader : _monad<String>, ρρFilename
	{
		//(0) [1 (3) ^ cl-cl_runon-cma_c aj-hd_scp_c av_-_s-cp-mc-pr_le "well"] -0.104715 {5845 187 5845 187} [0 1]
		//(1) [1 (2) cl-cl_runon-cma_c aj-hd_scp_c av_-_s-cp-mc-pr_le "well"] -0.1424 {6339 200 6339 200} [0 1]
		//(2) [1 (1) aj-hd_scp_c av_-_s-cp-mc-pr_le "well"] 0.31382 {35927 457 35925 455} [0 2]
		const string featurePattern = @"^\(\d+\) \[(?<featname>[^\]]+)\] (?<weight>\S+) {";
		static readonly Regex regex;

		static MaxEntModelLoader()
		{
			regex = new Regex(featurePattern, RegexOptions.Compiled);
		}

		public MaxEntModelLoader(IRtParent parent, String filename)
			: base(parent, filename)
		{
		}
		public MaxEntModelLoader(IRtParent parent)
			: this(parent, default(String))
		{
		}

		public String Filename
		{
			get { return base.t0; }
			set { base.t0 = value; }
		}

		public int m_grandparenting;

		public SlotDictionary<String, Double> dict;

		MaxEntModelLoader _load()
		{
			using (var mmf = MemoryMappedFile.CreateFromFile(_use_file, FileMode.Open))
			using (var vs = mmf.CreateViewStream())
			{
				dict = new SlotDictionary<String, Double>((int)(vs.Length / 100));

				using (StreamReader sr = new StreamReader(vs, Encoding.UTF8))
				{
					//////////////////////////////////////////////////////
					/// HEADER
					//////////////////////////////////////////////////////
					string line;
					do
					{
						if ((line = sr.ReadLine()) == null)
							throw new Exception(String.Format("Invalid maxent model file {0}. Expected ':begin :features", _use_file));
						if ((line = line.Trim().TrimEnd('.')).Length == 0 || line[0] == ';')
							continue;
						String[] rgs = line.Split(Charset.sp, StringSplitOptions.RemoveEmptyEntries);
						if (rgs.Length == 3 && rgs[1] == ":=")
						{
							switch (rgs[0])
							{
								case "*feature-grandparenting*":
									m_grandparenting = int.Parse(rgs[2]);
									break;
							}
						}
					}
					while (!line.StartsWith(":begin :features"));

					//////////////////////////////////////////////////////
					/// FEATURES
					//////////////////////////////////////////////////////
					while ((line = sr.ReadLine()) != null)
					{
						if ((line = line.Trim()).Length == 0)
							continue;

						Match m = regex.Match(line);
						if (m.Success)
							dict[m.Groups["featname"].Value] = Double.Parse(m.Groups["weight"].Value);
					}
				}
			}
			return this;
		}

		String _use_file;

		public override Object ProvideValue(IServiceProvider sp)
		{
			var fi = DirectoryPathSet.FindFileInfo(this, t0);
			if (fi == null)
				throw new FileNotFoundException("Maxent model file not found");

			this._use_file = fi.FullName;

			return Task.Factory.StartNew<MaxEntModelLoader>(_load);
		}
	};


	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	/// <summary>
	/// 
	/// </summary>
	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	[DisplayName("MaxEnt Model")]
	[Description("Maximum Entropy probablistic parse selection model, trained over hierarchical features in derivation trees.")]
	public sealed class MaxEntModel : identity_named
	{
		public MaxEntModel(IIdentity prv)
			: base(prv)
		{
		}

		public SlotDictionary<String, Double> dict;

		int m_grandparenting;
		public int Grandparenting
		{
			get { return this.m_grandparenting; }
			set { this.m_grandparenting = value; }
		}

		Task loading;
		public Task<MaxEntModelLoader> Loader
		{
			set
			{
				loading = value.ContinueWith(t =>
					{
						MaxEntModelLoader ldr = t.Result;
						this.dict = ldr.dict;
						this.m_grandparenting = ldr.m_grandparenting;
						Interlocked.Exchange(ref loading, null);
					});
			}
		}

		public struct ScoreDetail
		{
			public ScoreDetail(IDerivation node)
			{
				this.node = node;
				this.base_score = 0.0;
				this.dtrs = 0.0;
				this.gp = null;
			}
			public IDerivation node;
			public double base_score;
			public double dtrs;
			public double[] gp;

			public Double Total { get { return base_score + dtrs + GpTotal; } }
			public Double TotalExceptDtrs { get { return base_score + GpTotal; } }
			public Double GpTotal { get { return gp.Sum(); } }

			public String ToString(bool f_color)
			{
				String fmt;
				if (f_color)
					fmt = "$darkmagenta #{0,-4:X}$";
				else
					fmt = "#{0,-4:X}";
				fmt += " {1,-20} tot: {2,11:0.#######} base: {3,11:0.#######} dtrs: {4,11:0.#######} gp: {5,11:0.#######} ({6})";

				return String.Format(fmt,
					0,//node.Id,
					node.ChartLicense.Name,
					Total,
					base_score,
					dtrs,
					GpTotal,
					gp.StringJoin(" "));
			}
		};

		//choose the correct form for the name.  Use "type" for lexical entries, and "rule" for other nodes.
		static String getIDerivName(IDerivation deriv)
		{
			return deriv.Daughters.Count == 0 ? deriv.ChartLicense.InstanceType.Name : deriv.ChartLicense.Name;
		}

		///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
		/// <summary>
		/// TEMPLATE 1
		/// </summary>
		///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

		const String Template1 = "1 ({0}) {1}";

		public static IEnumerable<String> ExtractTemplate1Features(IDerivation deriv, int m_grandparenting)
		{
			return template1HelperS(deriv, new List<String> { "^" }, m_grandparenting);
		}
		public IEnumerable<Double> ExtractTemplate1FeatureValues(IDerivation deriv)
		{
			var _tmp = loading;
			if (_tmp != null)
				_tmp.Wait();
			return template1Helper(deriv, new List<String> { "^" });
		}
		public ScoreDetail ExtractTemplate1FeatureValuesExt(IDerivation deriv)
		{
			var _tmp = loading;
			if (_tmp != null)
				_tmp.Wait();
			return template1HelperExt(deriv, new List<String> { "^" });
		}

		///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
		/// <summary>
		/// TEMPLATE 1
		/// </summary>
		///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

		static IEnumerable<String> template1HelperS(IDerivation deriv, List<String> path, int m_grandparenting)
		{
			String curBase = null;

			//if this node is a lexical entry, use a special format for the edge (type rather than rule, and string)
			var le = deriv.ChartLicense as LexicalEntry;
			if (le != null)
			{
				curBase = String.Format("{0} \"{1}\"",
					le.InstanceType.Name,
					le.Text.ToLowerInvariant());
			}
			else
			{
				String curNode = curBase = deriv.ChartLicense.Name;

				//the "base" of the feature is the current node with the names of its daughters
				foreach (IDerivation dd in deriv.Daughters)
				{
					//recurse on daughters, adding the name of the current node to the path for grandparenting
					path.Add(curNode);
					foreach (String s in template1HelperS(dd, path, m_grandparenting))
						yield return s;
					curBase += String.Format(" {0}", getIDerivName(dd));
				}
			}

			// base feature with no grandparenting

			yield return String.Format(Template1, "0", curBase);

			// features with varying degrees of grandparenting
			for (int i = 0; i < Math.Min(path.Count, m_grandparenting); i++)
			{
				curBase = String.Format("{0} {1}", path[path.Count - 1 - i], curBase);
				yield return String.Format(Template1, (i + 1).ToString(), curBase);
			}

			//remove the last node in the "path" before returning (because the same path object is used throughout the tree)
			path.RemoveAt(path.Count - 1);
		}

		///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
		/// <summary>
		/// TEMPLATE 1
		/// </summary>
		///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
		IEnumerable<Double> template1Helper(IDerivation deriv, List<String> path)
		{
			double v;
			foreach (var s in template1HelperS(deriv, path, m_grandparenting))
				if (dict.TryGetValue(s, out v))
					yield return v;
		}

		///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
		/// <summary>
		/// TEMPLATE 1
		/// </summary>
		///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

		ScoreDetail template1HelperExt(IDerivation deriv, List<String> path)
		{
			String curBase = null;

			ScoreDetail sd = new ScoreDetail(deriv);

			//if this node is a lexical entry, use a special format for the edge (type rather than rule, and string)
			var le = deriv.ChartLicense as LexicalEntry;
			if (le != null)
			{
				curBase = String.Format("{0} \"{1}\"",
					le.InstanceType.Name,
					le.Text.ToLowerInvariant());
			}
			else
			{
				String curNode = curBase = deriv.ChartLicense.Name;

				//the "base" of the feature is the current node with the names of its daughters
				foreach (IDerivation dd in deriv.Daughters)
				{
					//recurse on daughters, adding the name of the current node to the path for grandparenting
					path.Add(curNode);
					sd.dtrs += template1HelperExt(dd, path).Total;
					curBase += String.Format(" {0}", getIDerivName(dd));
				}
			}

			// base feature with no grandparenting
			double v;
			if (dict.TryGetValue(String.Format(Template1, "0", curBase), out v))
				sd.base_score = v;

			// features with varying degrees of grandparenting
			sd.gp = new double[Math.Min(path.Count, m_grandparenting)];
			for (int i = 0; i < sd.gp.Length; i++)
			{
				curBase = String.Format("{0} {1}", path[path.Count - 1 - i], curBase);
				if (dict.TryGetValue(String.Format(Template1, (i + 1).ToString(), curBase), out v))
					sd.gp[i] = v;
			}

			//remove the last node in the "path" before returning (because the same path object is used throughout the tree)
			path.RemoveAt(path.Count - 1);

			return sd;
		}
	};

	public struct MaxentFeatureSet : IIdentity, ρρParseSelectionFeatures
	{
		public IDerivation deriv;

		public String features;
		public String ParseSelectionFeatures
		{
			get { return features; }
			set { features = value; }
		}

		IIdentity IIdentity.Trace { get { return deriv; } }
	};

	public class MaxentFeatureList : IdentArray<MaxentFeatureSet>
	{
		public MaxentFeatureList(IIdentity prv, MaxentFeatureSet[] arr)
			: base(prv, arr)
		{
		}
	};


	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	/// 
	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	public sealed class DerivationFeatureExtractor : m_scalar_base<IIdentList<IDerivation>, MaxentFeatureList>
	{
		public DerivationFeatureExtractor(IRtParent parent, IIdentList<IDerivation> item)
			: base(parent, item)
		{
		}

		protected override MaxentFeatureList start(IIdentList<IDerivation> rgd)
		{
			var arr = new MaxentFeatureSet[rgd.Count];
			StringBuilder sb = new StringBuilder();

			for (int i = 0; i < rgd.Count; i++)
			{
				sb.Clear();

				var e = MaxEntModel.ExtractTemplate1Features(rgd[i], 3).GetEnumerator();
				while (e.MoveNext())
					sb.AppendLine(e.Current);

				arr[i].deriv = rgd[i];
				arr[i].features = sb.ToString();
			}
			return new MaxentFeatureList(rgd, arr);
		}
	};


	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	/// actually, all this monad does is put the model onto the derivation, enabling it to
	/// lazily score itself later
	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	public sealed class ScorerLoader : m_scalar_base<IDerivation, IDerivation>,
		ρρMaxEntModel
	{
		public ScorerLoader(IRtParent parent, IDerivation item)
			: base(parent, item)
		{
		}

		MaxEntModel _model;
		public MaxEntModel MaxEntModel
		{
			get { return _model; }
			set { _model = value; }
		}

		protected override IDerivation start(IDerivation deriv)
		{
			P.MaxEntModel.SetValue(deriv, _model);
			return deriv;
		}
	};

	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	/// 
	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	public sealed class BatchScorer : m_scalar_base<IIdentList<IDerivation>, IIdentList<IDerivation>>,
		IComparer<IDerivation>,
		ρρMaxEntModel
	{
		public BatchScorer(IRtParent parent, IIdentList<IDerivation> item)
			: base(parent, item)
		{
			this.TopN = int.MaxValue;
			this.SortEvenIfFewerThanN = true;
		}

		public int TopN { get; set; }

		public bool SortEvenIfFewerThanN { get; set; }

		MaxEntModel _model;
		public MaxEntModel MaxEntModel
		{
			get { return _model; }
			set { _model = value; }
		}

		protected override IIdentList<IDerivation> start(IIdentList<IDerivation> t)
		{
			if (_model == null)
				throw new RequiredPropertyNotSetException("MaxEntModel");

			int n = TopN;
			if (n < 0 || n > t.Count)
				n = t.Count;

			IDerivation[] best = null;
			if (n > 0)
			{
				best = new IDerivation[n];
				if (n == t.Count)
					score_all(t, best);
				else
					score_best(t, best);
			}
			return new IdentArray<IDerivation>(this, best);
		}

		void score_all(IIdentList<IDerivation> t, IDerivation[] best)
		{
			int ix = 0;
			foreach (var deriv in t)
				(best[ix++] = deriv).Score(_model);

			if (SortEvenIfFewerThanN)
				Array.Sort<IDerivation>(best, this);
		}

		void score_best(IIdentList<IDerivation> t, IDerivation[] best)
		{
			int c_items = 0;
			int cx = best.Length - 1;

			foreach (var deriv in t)
			{
				var score = deriv.Score(_model);

				if (c_items <= cx)
					best[c_items] = deriv;
				else
				{
					if (c_items == best.Length)		/// one-time sort when 'N' items are reached
						Array.Sort<IDerivation>(best, this);

					if (score >= retrieve_score(best[cx]))
					{
						int ix = 0;
						if (score <= retrieve_score(best[0]))
						{
							ix = Array.BinarySearch<IDerivation>(best, deriv, this);
							if (ix < 0)
								ix = ~ix;
						}

						if (ix < cx)
							Array.Copy(best, ix, best, ix + 1, cx - ix);
						best[ix] = deriv;
					}
				}
				c_items++;
			}

#if true
			Debug.Assert(c_items > best.Length);
#else
			if (c_items < best.Length)
			{
				best = alib.Array.arr.Resize(best, c_items);

				if (SortEvenIfFewerThanN)
					Array.Sort<IDerivation>(best, this);
			}
#endif
		}

		int IComparer<IDerivation>.Compare(IDerivation x, IDerivation y)
		{
			double d1 = retrieve_score(x);
			double d2 = retrieve_score(y);
			return d1 == d2 ? 0 : d1 < d2 ? 1 : -1;
		}

		static Double retrieve_score(IDerivation x)
		{
			return (Double)P.ParseSelectionScore.GetValue(x);
		}
	};
}
